#include "Rabin.h"

#include <vector>
#include <stdexcept>
#include <ctime>

#include <gmp.h>

#include "../DEBUG_PRINT_MACRO.h"


#include "data_conversion.h"

#ifdef DEBUG_PRINTS_LEVEL_1
#include <iostream>
#endif

// 引入 Tracy 头文件
#include <Tracy.hpp>


template<int key_length>
void Rabin<key_length>::generatepublicparameters(Rabin_publicparameter& publicParameters) {
    // Rabin算法通常不需要额外公共参数，留空实现
    ZoneScopedN("Rabin::generatepublicparameters"); // 标记整个函数
}

template<int key_length>
void Rabin<key_length>::generateKeyPair(Rabin_privatekey& privatekey, Rabin_publickey& publickey, Rabin_publicparameter& publicparameter) {
    ZoneScopedN("Rabin::generateKeyPair"); // 标记整个函数

    mpz_t p, q, n, temp;
    {
        ZoneScopedN("Rabin::generateKeyPair::GMP Variable Initialization");
        mpz_inits(p, q, n, temp, nullptr);
    }
    gmp_randstate_t state;
    {
        ZoneScopedN("Rabin::generateKeyPair::Random State Initialization");
        gmp_randinit_default(state);
        gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
    }


    int L = key_length;

    const int MAX_P_ATTEMPTS = 1000;
    const int MAX_Q_ATTEMPTS = 1000;
    int attempts_p = 0;
    int attempts_q = 0;


    // --------------------------
    // 1. 生成素数p满足p ≡ 3 mod 4
    // --------------------------
    {
        ZoneScopedN("Rabin::generateKeyPair::Generate Prime P");
        do {
            mpz_urandomb(p, state, L / 2);
            mpz_nextprime(p, p);
        } while (mpz_mod_ui(temp, p, 4) != 3 || ++attempts_p > MAX_P_ATTEMPTS);
    }

    // --------------------------
    // 2. 生成素数q满足q ≡ 3 mod 4
    // --------------------------
    {
        ZoneScopedN("Rabin::generateKeyPair::Generate Prime Q");
        do {
            mpz_urandomb(q, state, L / 2);
            mpz_nextprime(q, q);
        } while (mpz_mod_ui(temp, q, 4) != 3 || ++attempts_q > MAX_Q_ATTEMPTS);
    }

    if (attempts_p > MAX_P_ATTEMPTS || attempts_q > MAX_Q_ATTEMPTS) {
        throw std::runtime_error("Failed to generate suitable prime numbers.");
    }

    // --------------------
    // 3. 计算公钥n = p * q
    // --------------------
    {
        ZoneScopedN("Rabin::generateKeyPair::Calculate Public Key N");
        mpz_mul(n, p, q);
    }

    {
        ZoneScopedN("Rabin::generateKeyPair::Store Keys and Cleanup");
        privatekey.p = GMPInt(p);
        privatekey.q = GMPInt(q);
        publickey.n = GMPInt(n);

        mpz_clears(p, q, n, temp, nullptr);
        gmp_randclear(state);
    }
}

template<int key_length>
Rabin_signature Rabin<key_length>::sign(const PlainMessage& message, const Rabin_privatekey& privatekey, const Rabin_publicparameter& publicparameter) {
    ZoneScopedN("Rabin::sign"); // 标记整个函数

#ifdef DEBUG_PRINTS_LEVEL_3
    std::cout << "[Rabin sign][debug]Starting Rabin signing process..." << std::endl;
#endif

    mpz_t p, q, n, hash_m, u, temp, inv, s_p, s_q, exp, s;
    {
        ZoneScopedN("Rabin::sign::GMP Variable Initialization");
        mpz_inits(p, q, n, hash_m, u, temp, inv, s_p, s_q, exp, s, nullptr);
        mpz_set(p, privatekey.p);
        mpz_set(q, privatekey.q);
        mpz_mul(n, p, q);
    }

    gmp_randstate_t state;
    {
        ZoneScopedN("Rabin::sign::Random State Initialization");
        gmp_randinit_default(state);
        gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
    }

    Rabin_signature sig;
    bool found = false;

    while (!found) {
        ZoneScopedN("Rabin::sign::Quadratic Residue Search Loop"); // 标记寻找二次剩余的循环
        // ---------------------
        // 1. 生成随机数u（128位）
        // ---------------------
        {
            ZoneScopedN("Rabin::sign::Generate Random U");
            mpz_urandomb(u, state, 128);
        }
#ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[Rabin sign][debug]: u = %Zd\n", u);
#endif
        // 将u转换为字节流
        size_t u_size;
        {
            ZoneScopedN("Rabin::sign::Convert U to Bytes");
            u_size = (mpz_sizeinbase(u, 2) + 7) / 8;
            std::vector<uint8_t> u_bytes(u_size);
            mpz_export(u_bytes.data(), nullptr, 1, 1, 0, 0, u); // 大端序

            // ----------------------
            // 2. 计算hash_m = H(m||u)
            // ----------------------
            HASH_VALUE combined_hash = message.get_hash_data_with_suffix(u_bytes);

            // 将哈希值转换为mpz整数
            mpz_import(hash_m, 32, 1, 1, 0, 0, combined_hash.data());
        }
#ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[Rabin sign][debug]: hash_m = %Zd\n", hash_m);
#endif
        // 检查二次剩余性质
        {
            ZoneScopedN("Rabin::sign::Check Quadratic Residue");
            int legendre_p = mpz_legendre(hash_m, p);
            int legendre_q = mpz_legendre(hash_m, q);
            if (legendre_p == 1 && legendre_q == 1) found = true;
        }
    }

    // ---------------------------------------------
    // 3. 计算hash_m的平方根s_p,s_q（保持原CRT部分不变）
    // ---------------------------------------------
    {
        ZoneScopedN("Rabin::sign::Calculate Square Roots (s_p, s_q)");
        mpz_add_ui(exp, p, 1);
        mpz_divexact_ui(exp, exp, 4);
        mpz_powm(s_p, hash_m, exp, p);  // s_p = c ^ ((p + 1) / 4) mod p
    }
#ifdef DEBUG_PRINTS_LEVEL_3
    gmp_printf("[Rabin sign][debug]: s_p = %Zd\n", s_p);
#endif
    {
        ZoneScopedN("Rabin::sign::Calculate Square Roots (s_q)");
        mpz_add_ui(exp, q, 1);
        mpz_divexact_ui(exp, exp, 4);
        mpz_powm(s_q, hash_m, exp, q);  // s_q = c ^ ((q + 1) / 4) mod q
    }
#ifdef DEBUG_PRINTS_LEVEL_3
    gmp_printf("[Rabin sign][debug]: s_q = %Zd\n", s_q);
#endif

    // ---------------------------------------------
    // 4. 找到s满足： s ≡ s_p mod p && s ≡ s_q mod q
    // ---------------------------------------------
    {
        ZoneScopedN("Rabin::sign::Apply CRT to Find S");
        mpz_sub(temp, s_q, s_p);    // t = s_q - s_p
        mpz_invert(inv, p, q);
        mpz_mul(temp, temp, inv);
        mpz_mod(temp, temp, q);
        mpz_mul(temp, temp, p);
        mpz_add(s, s_p, temp);
    }

    // 设置签名
    {
        ZoneScopedN("Rabin::sign::Store Signature and Cleanup");
        sig.s = GMPInt(s);
        sig.u = GMPInt(u);

        // 清理资源
        mpz_clears(p, q, n, hash_m, u, temp, inv, s_p, s_q, exp, s, nullptr);
        gmp_randclear(state);
    }
    return sig;
}



template<int key_length>
bool Rabin<key_length>::verify(const PlainMessage& message, const Rabin_publickey& publickey, const Rabin_publicparameter& publicparameter, const Rabin_signature& signature) {
    ZoneScopedN("Rabin::verify"); // 标记整个函数

    mpz_t n, hash_m, s_sq;
    {
        ZoneScopedN("Rabin::verify::GMP Variable Initialization");
        mpz_inits(n, hash_m, s_sq, nullptr);
        mpz_set(n, publickey.n);
    }

    // 将u转换为字节流
    size_t u_size;
    std::vector<uint8_t> u_bytes;
    {
        ZoneScopedN("Rabin::verify::Convert U to Bytes");
        u_size = (mpz_sizeinbase(signature.u, 2) + 7) / 8;
        u_bytes.resize(u_size);
        mpz_export(u_bytes.data(), nullptr, 1, 1, 0, 0, signature.u); // 大端序
    }
    // -------------
    // 1. 计算H(m||u)
    // -------------
    HASH_VALUE combined_hash;
    {
        ZoneScopedN("Rabin::verify::Calculate H(m||u)");
        combined_hash = message.get_hash_data_with_suffix(u_bytes);

        // 将哈希值H(m||u)转换为mpz整数
        mpz_import(hash_m, 32, 1, 1, 0, 0, combined_hash.data());
    }

    // -----------------
    // 2. 计算 s_sq = s^2 mod n
    // -----------------
    {
        ZoneScopedN("Rabin::verify::Calculate s_sq");
        mpz_powm_ui(s_sq, signature.s, 2, n);
    }
    bool valid;
    {
        ZoneScopedN("Rabin::verify::Compare Hash and s_sq");
        valid = (mpz_cmp(hash_m, s_sq) == 0); // 判断 s_sq 与 hash_m 是否相等
    }

#ifdef DEBUG_PRINTS_LEVEL_3
    gmp_printf("[Rabin sign][debug]: hash_m = %Zd\n", hash_m);
    gmp_printf("[Rabin sign][debug]: s_sq = %Zd\n", s_sq);
#endif

    {
        ZoneScopedN("Rabin::verify::Cleanup Resources");
        mpz_clears(n, hash_m, s_sq, nullptr);
    }
    return valid;
}

template class Rabin<1024>;
template class Rabin<2048>;
template class Rabin<3072>;